package mantel;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.Random;

import tool.Utility;
import jsc.combinatorics.Permutation;
import jsc.combinatorics.Permutations;
import jsc.correlation.SpearmanCorrelation;
import jsc.datastructures.PairedData;
import jsc.regression.PearsonCorrelation;

public class mantelTest {
	
	public static double[][] mantel_cor(double[][] X, double[][] Y, int permutations, int breaks,boolean standardise){
		double[][] mstat = new double[breaks+1][2];
		
		double mX = max(Y);
		double interval = mX/breaks;
		double old = 0.0;
		for (int i = 0; i < breaks+1; i++) {
			double[] bounds = {old,old+interval};
			double[][] model = getModelMatrix(Y, bounds);
			double[] mVal = mantel(X,model,permutations,standardise);
			double[] statline = {old,mVal[0],mVal[1]};
			mstat[i] = statline; 
			old = old+interval;
		}
		return mstat;
	}
	
	public static double max(double[][] X){
		//obtains max from the upper triangle of a matrix
		double max = 0.0;
		for (int i = 1; i < X.length; i++) {
			for (int j = 0; j < i; j++) {
				max = (X[i][j]>max) ? X[i][j] : max;
			}
		}
		return max;
	}
	
	public static void printCorrelog(double[][] cor){
		for(int i=0;i<cor.length;i++){
			System.out.println(cor[i][0]+" Correlation: "+cor[i][1]+" Significance: "+cor[i][2]);
		}
	}
	
	
	public static ArrayList<Double> getBreakPoints(double[][] X, int groupSize){
		
		ArrayList<Double> breakPoints = new ArrayList<Double>();
		
		double[] vals = Utility.distMatrixToVector(X);
		Arrays.sort(vals);
		double current = vals[0];
		breakPoints.add(current);
		int groupNum = 0;
		for (int i = 0; i < vals.length; i++) {
			if(Utility.roundToSignificantFigures(vals[i],5)!=current){
				if(groupNum >=groupSize){
					breakPoints.add(current);
					groupNum = 1;
				}
				else {
					groupNum = groupNum+1;
				}
			}
			current = Utility.roundToSignificantFigures(vals[i], 5);
		}
		return breakPoints;
	}
	
	public static double[][] segmented_mantel_cor(double[][] X, double[][] Y, int permutations, int groupSize){
		
		double mx = max(X);
		
		ArrayList<Double> breakPoints = getBreakPoints(X,groupSize);
		
		double[][] mstat = new double[breakPoints.size()][3];
		
		System.out.println("Using breakpoints "+breakPoints);

		int i=0;
		for(Double dLower:breakPoints){
		
			Double dUpper;
			if(i==breakPoints.size()-1){
				dUpper = mx+0.1; //hack to ensure final value accounted for.
			}
			else {
				dUpper = breakPoints.get(i+1);
			}
			
			double[] bounds = {dLower,dUpper};
			System.out.println("Calculating for bounds: "+bounds[0]+" "+bounds[1]);
			
			double[] mVal;
			try {
				mVal = bounded_mantel(X,Y,permutations,bounds);
				double[] statline = {dLower,mVal[0],mVal[1]};
				mstat[i] = statline;
			}
			catch(NullPointerException e){
				System.out.println("Issue with distance catagory.");
				mstat[i][0] = dLower;
			}
			i = i+1;
		}
		return mstat;
	}
	
	public static double[] bounded_mantel(double[][] X, double[][] Y, int permutations, double[] bounds){
		double[] XV = Utility.distMatrixToVector(X);
		double[] YV = Utility.distMatrixToVector(Y);
		double[][] rVec_original = getReducedVectors(XV,YV,bounds);
		System.out.println("Elements in distance class: "+rVec_original[0].length);
	
		
		/*for(int i=0;i<rVec_original[0].length;i++){
			System.out.println(rVec_original[0][i]+" "+rVec_original[1][i]);
		}
		System.out.println("----");*/
		
		//double[] c = pcor(rVec_original[0],rVec_original[1]);
		double[] c = cor(rVec_original[0],rVec_original[1]);
		if(c==null){
			System.out.println("Empty distance catagory: "+bounds[0]+" "+bounds[1]);
			return null;
		}
		
		Permutations permuter = new Permutations(X.length);
	
		double sig = 0;
		for (int i = 0; i < permutations; i++) {
			//System.out.println("Permutations: "+i);
			double[][] XP = permuteMatrix(X,permuter);
			double[] XPV = Utility.distMatrixToVector(XP);
			double[][] rVec = getReducedVectors(XPV,YV,bounds);
			double[] cPerm = cor(rVec[0],rVec[1]);
			if(c[0]>0){
				sig = (cPerm[0]>=c[0]) ? sig+1 : sig;
			}
			else {
				sig = (cPerm[0]<=c[0]) ? sig+1 : sig;
			}
		}
		c[1] = (sig+1)/(permutations+1);
		
		return c;
	}
	
	private static double[][] getReducedVectors(double[] X, double[] Y, double[] bounds){
		double lowerbound = bounds[0];
		double upperbound = bounds[1];
		
		double[] X_Reduced = new double[X.length];
		double[] Y_Reduced = new double[Y.length];
		
		int j = 0;
		for(int i=0;i<X.length;i++){
			double xi = X[i];
			if(lowerbound <= xi && upperbound>xi){
				X_Reduced[j]=(xi);
				Y_Reduced[j]=Y[i];
				j=j+1;
			}
		}
		
		if(j==0){System.out.println("Warning: no data within bounds for reduced vector.");}
		
		double[][] rvec = new double[2][j+1];
		rvec[0] = Arrays.copyOf(X_Reduced, j);
		rvec[1] = Arrays.copyOf(Y_Reduced, j);
		return rvec;
	}
	
	public static double[] mantel(double[][] X, double[][] Y, int permutations, boolean standardise){
		if(standardise){
			X = standardise(X);
			Y = standardise(Y);
		}
		
		double[] XV = Utility.distMatrixToVector(X);
		double[] YV = Utility.distMatrixToVector(Y);
		double[] c = pcor(XV,YV);
		
		Permutations permuter = new Permutations(X.length);
	
		double sig = 0;
		for (int i = 0; i < permutations; i++) {
			//System.out.println("Permutations: "+i);
			double[][] XP = permuteMatrix(X,permuter);
			double[] XPV = Utility.distMatrixToVector(XP);
			double[] cPerm = pcor(XPV,YV);
			if(c[0]>0){
				sig = (cPerm[0]>=c[0]) ? sig+1 : sig;
			}
			else {
				sig = (cPerm[0]<=c[0]) ? sig+1 : sig;
			}
		}
		c[1] = (sig+1)/(permutations+1);
		
		return c;
	}
	
	public static double[][] standardise(double[][] X){
		double[][] Y = new double[X.length][X.length];
		double[] s = stat(X);
		for (int i = 0; i < X.length; i++) {
			for (int j = 0; j < X.length; j++) {
				Y[i][j] = (X[i][j]-s[0])/s[1];
			}
		}
		return Y;
	}
	
	public static double[] stat(double[][] X){
		double mean = 0.0;
		double sqmean = 0.0;
		for (int i = 0; i < X.length; i++) {
			double[] xr = X[i];
			double[] xr2 = squares(xr);
			mean = mean+mean(xr);
			sqmean = sqmean+mean(xr2);
		}
		mean = mean/X.length;
		sqmean = sqmean/X.length;
		double stdev = Math.sqrt(sqmean-(mean*mean));
		double[] s = {mean,stdev};
		return s;
	}
	
	public static double[] squares(double[] x){
		double[] y = new double[x.length];
		for (int i = 0; i < x.length; i++) {
			double s = x[i];
			y[i] = s*s;
		}
		return y;
	}
	
	public static double[][] permuteMatrix(double[][] X, Permutations permuter){
		Permutation p = permuter.randomPermutation();
		int[] k = p.toIntArray();
		int[] k0 = new int[k.length];
		for(int i=0;i<k.length;i++){k0[i] = k[i]-1;}
		double[][] XP = new double[X.length][X.length];
		for(int i=0;i<k.length;i++){
			double[] rowX = X[k0[i]];
			double[] rowXP = new double[rowX.length];
			for(int j=0;j<k.length;j++){
				rowXP[j]= rowX[k0[j]];
			}
			XP[i] = rowXP;
		}
		return XP;
	}
	
	public static double[] cor(double[] x, double[] y){
		double[] c = new double[2];
		try{
			if(x==null || y==null || x.length==0 || y.length==0){
				throw new IllegalArgumentException();
			}
		}
		catch (IllegalArgumentException e) {
			System.out.println("Illegal argument to cor");
			return null;
		}
		
		
		PairedData pxy = new PairedData(x,y);
		PearsonCorrelation spc;
		
		//System.out.println(pxy);
		try{
			spc = new PearsonCorrelation(pxy);
		}
		catch(IllegalArgumentException e){
			System.out.println("Constant data, checking for equality...");
			if(Arrays.equals(x, y)){
				c[0] = 1.0; return c;
			}
			else {
				System.out.println("Data is constant but not equal: need to respecify distance catagories.");
				spc=null;
				c[0] = 0.0;
				return null;
			}
		}
		c[0] = spc.getR();
		return c;
	}
	
	public static double[] pcor(double[] x, double[] y){
		//computes sample correlation coefficient - faster than cor
		//because it does not calculate significance / convert to paired data.
		//Will return NAN if stdev =0
		if(Arrays.equals(x, y)){
			double[] c = {1.0,0.0};
			return c;
		}
		
		double muX = mean(x); double muY = mean(y);
		double U = 0.0;
		double L1 = 0.0;
		double L2 = 0.0;
		
		for(int i=0;i<x.length;i++){
			double xD = x[i]-muX; double yD = y[i]-muY;
			double xD2 = xD*xD; double yD2 = yD*yD;
			U = U+(xD*yD);
			L1 = L1+xD2; L2 = L2+yD2;
		}
		
		double[] c ={U/Math.sqrt(L1*L2),0.0}; //does not calculate sig
		//returned with redundant value to conform with cor.
		return c;
		
	}
	
	public static double mean(double[] x){
		double sm = 0.0;
		for(int i=0;i<x.length;i++){
			sm = sm+x[i];
		}
		return sm/(x.length);
	}
	
	public static void main(String args[]){
		
		
		double[][] X = {{0.0,1.0,2.0,3.0},{1.0,0.0,3.0,4.0},{2.0,3.0,0.0,4.0},{3.0,4.0,0.0,0.0}};
		//double[][] Y = {{0.0,1.0,2.0,3.0},{1.0,0.0,3.0,4.0},{2.0,3.0,0.0,4.0},{3.0,4.0,0.0,0.0}};
		
		Utility.printMat(X);
		System.out.println(max(X));
		
		/*Random r = new Random();
		double[][] X = new double[20][20];
		double[][] Y = new double[20][20];
		for(int i=0;i<20;i++){
			for(int j=0;j<20;j++){
				X[i][j] = r.nextInt(21);
				if(X[i][j]<=3){
					Y[i][j] = X[i][j];
				}
				else {
					Y[i][j] = r.nextInt(21);
				}
			}
		}
		
		Utility.printMat(X);
		System.out.println();
		Utility.printMat(Y);*/
		
		
		//long t1 = System.currentTimeMillis();
		//double[] test = mantel(X,Y,1000,false);
		//long t2 = System.currentTimeMillis();
		//System.out.println("Correlation: "+test[0]+" Sig: "+test[1]);
		//System.out.println("Total Time: "+(t2-t1));
		//System.out.println();
		
		//double[][] mcor = segmented_mantel_cor(X,Y,1000,5,20);
		//printCorrelog(mcor);
		
	}
	
	public static double[][] getModelMatrix(double[][] X, double[] bounds){
		double lowerbound = bounds[0];
		double upperbound = bounds[1];
		
		double[][] Y = new double[X.length][X.length];
		
		for (int i = 0; i < X.length; i++) {
			for (int j = 0; j < X.length; j++) {
				double xij = X[i][j];
				Y[i][j] = (lowerbound<=xij && xij<upperbound)? 0:1;  
			}
		}
		return Y;
	}
}
